package org.dfzt.config.filter;

import org.apache.commons.lang3.StringUtils;
import org.dfzt.common.util.XssFilterUtil;
import org.dfzt.common.util.encryption.SMUtils;

import javax.servlet.ReadListener;
import javax.servlet.ServletInputStream;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletRequestWrapper;
import java.io.BufferedReader;
import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.io.InputStreamReader;
import java.nio.charset.StandardCharsets;
import java.util.Map;

/**
 * 防止xss攻击，包装请求与响应的包装类
 * @author furao
 */
public class XssHttpServletRequestWrapper extends HttpServletRequestWrapper {
    HttpServletRequest orgRequest = null;
    private byte[] body;
    private boolean isIncludeRichText = false;
    private Map<String, String[]> paramMap=null;
    private boolean isEncrypt=false;

    public XssHttpServletRequestWrapper(HttpServletRequest request, boolean isIncludeRichText, boolean isEncrypt) {
        super(request);
        this.isEncrypt=isEncrypt;
        orgRequest = request;
        String method=request.getMethod();
        if(this.isEncrypt){
            if("get".equalsIgnoreCase(method)){
                paramMap=HttpHelper.getParamMap(request);
            }else{
                body = HttpHelper.getBodyString(request).getBytes(StandardCharsets.UTF_8);
            }
        }
        this.isIncludeRichText = isIncludeRichText;
    }

    /**
     * 覆盖getParameter方法，将参数名和参数值都做xss过滤.
     * 如果需要获得原始的值，则通过super.getParameterValues(name)来获取
     * getParameterNames,getParameterValues和getParameterMap也可能需要覆盖
     */
    @Override
    public String getParameter(String name) {
        if (("content".equals(name) || name.endsWith("WithHtml")) && !isIncludeRichText) {
            return super.getParameter(name);
        }
        name = XssFilterUtil.clean(name);
        String value = super.getParameter(name);
        if(this.isEncrypt){
            //解密
            value= SMUtils.sm2Decrypt(value);
        }
        //xss
        return XssFilterUtil.clean(value);
    }

    @Override
    public String[] getParameterValues(String name) {
        String[] arr = super.getParameterValues(name);
        if (arr != null) {
            for (int i = 0; i < arr.length; i++) {
                String value=arr[i];
                if(this.isEncrypt){
                    //解密
                    value= SMUtils.sm2Decrypt(arr[i]);
                }
                arr[i] = XssFilterUtil.clean(value);
            }
        }
        return arr;
    }

    @Override
    public Map<String, String[]> getParameterMap() {
        if(this.isEncrypt){
            //解密
            return paramMap;
        }else{
            return super.getParameterMap();
        }
    }

    /**
     * 覆盖getHeader方法，将参数名和参数值都做xss过滤。<br/>
     * 如果需要获得原始的值，则通过super.getHeaders(name)来获取<br/>
     * getHeaderNames 也可能需要覆盖
     */
    @Override
    public String getHeader(String name) {
        name = XssFilterUtil.clean(name);
        String value = super.getHeader(name);
        if (StringUtils.isNotBlank(value)) {
            value = XssFilterUtil.clean(value);
        }
        return value;
    }

    @Override
    public BufferedReader getReader(){
        return new BufferedReader(new InputStreamReader(getInputStream()));
    }

    @Override
    public ServletInputStream getInputStream(){

        if(!this.isEncrypt){
            try {
                return super.getInputStream();
            } catch (IOException e) {
                return null;
            }
        }
        final ByteArrayInputStream bais = new ByteArrayInputStream(body);

        return new ServletInputStream() {

            @Override
            public int read() throws IOException {
                return bais.read();
            }

            @Override
            public boolean isFinished() {
                return false;
            }

            @Override
            public boolean isReady() {
                return false;
            }

            @Override
            public void setReadListener(ReadListener listener) {

            }

        };
    }

}
